// Copyright 2011 Abhishek Kulkarni. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	"go/token"
	"go/ast"
)

const (
	propnet = "gopropnet"
)

var cells []*ast.Ident

func getCells(d ast.DeclStmt) ([]*ast.Ident, bool) {
	switch v := d.Decl.(type) {
	case *ast.GenDecl:
		if v.Tok == token.VAR {
			for i := range v.Specs {
				spec := v.Specs[i].(*ast.ValueSpec)
				if isName(spec.Type, "Cell") {
					return spec.Names, true
				}
			}
		}
	}
	return nil, false
}

func existsCell(c *ast.Ident) bool {
	for i := range(cells) {
		if c.Name == cells[i].Name {
			return true
		}
	}
	return false
}

func getPropagator(t token.Token) (*ast.Ident, bool) {
	switch t {
	case token.ADD:
		return ast.NewIdent("Sum"), true
	case token.SUB:
		return ast.NewIdent("Difference"), true
	case token.MUL:
		return ast.NewIdent("Product"), true
	case token.QUO:
		return ast.NewIdent("Division"), true
	}
	return nil, false
}

func transformPropagator(s ast.Stmt) ast.Stmt {
	switch d := s.(type) {
	case *ast.SendStmt:
		switch b := d.Value.(type) {
		case *ast.BinaryExpr:
			op, ok := getPropagator(b.Op)
			if (!ok) {
				return s
			}
			return &ast.ExprStmt{
			X: &ast.CallExpr{
				Fun:  &ast.SelectorExpr{ast.NewIdent(propnet), op},
				Args: []ast.Expr{b.X, b.Y, d.Chan},
				},
			}
		case *ast.BasicLit:
			return &ast.ExprStmt{
			X: &ast.CallExpr{
				Fun:  &ast.SelectorExpr{ast.NewIdent(propnet),
						ast.NewIdent("Constant")},
				Args: []ast.Expr{b, d.Chan},
				},
			}
		}
	}
	return s
}

func transformAST(f *ast.File) (fixed bool) {
	walkBeforeAfter(f, func(n interface{}) {
		b, ok := n.(*ast.ValueSpec)
		if !ok {
			return
		}
		cells = append(cells, b.Names...)
	}, func(n interface{}) {
		c, ok := n.(*ast.Stmt)
		if !ok {
			return
		}
		switch d := (*c).(type) {
		case *ast.AssignStmt:
			for i := range d.Rhs {
				id, ok := d.Rhs[i].(*ast.Ident)
				if !ok {
					continue
				}
				if existsCell(id) {
					d.Rhs[i] = &ast.CallExpr{
					Fun:  &ast.SelectorExpr{d.Rhs[i], ast.NewIdent("Content")},
					}
				}
			}

			id, ok := d.Lhs[0].(*ast.Ident)
			if (!ok) {
				return
			}
			if existsCell(id) {
				*c = &ast.ExprStmt{X: &ast.CallExpr{
					Fun:  &ast.SelectorExpr{d.Lhs[0], ast.NewIdent("Content")},
					Args: d.Rhs,
					},
				}
			}
		case *ast.ExprStmt:
			switch s := d.X.(type) {
			case *ast.CallExpr:
				for i := range s.Args {
					id, ok := s.Args[i].(*ast.Ident)
					if !ok {
						continue
					}
					if existsCell(id) {
						s.Args[i] = &ast.CallExpr{
						Fun:  &ast.SelectorExpr{s.Args[i], ast.NewIdent("Content")},
						}
					}
				}
			}
		}
	})

	walkBeforeAfter(f, func(n interface{}) {
		var newCells []ast.Stmt
		b, ok := n.(*ast.BlockStmt)
		if !ok {
			return
		}

		start := &ast.AssignStmt{
		Lhs: []ast.Expr{ast.NewIdent(propnet)},
		Tok: token.DEFINE,
		Rhs: []ast.Expr{
				&ast.CallExpr{
				Fun:  &ast.SelectorExpr{ast.NewIdent("goprop"),
						ast.NewIdent("NewPropagatorNetwork")},
				Args: []ast.Expr{&ast.BasicLit{Kind: token.STRING, Value: "\"" + propnet + "\""}},
				},
			},
		}

		var list []ast.Stmt
		list = append(list, start)
		list = append(list, b.List[0:]...)
		b.List = list

		for i := range b.List {
			switch d := b.List[i].(type) {
			case *ast.DeclStmt:
				cells, ok := getCells(*d)
				if !ok {
					continue
				}

				for j := range cells {
					cname := &ast.BasicLit{Kind: token.STRING, Value: "\"" + cells[j].Name + "\""}
					assignment := &ast.AssignStmt{
					Lhs: []ast.Expr{cells[j]},
					Tok: token.DEFINE,
					Rhs: []ast.Expr{
							&ast.CallExpr{
							Fun:  &ast.SelectorExpr{ast.NewIdent(propnet),
									ast.NewIdent("NewCell")},
							Args: []ast.Expr{cname},
							},
						},
					}
					newCells = append(newCells, assignment)
				}

				if len(newCells) > 0 {
					var list []ast.Stmt
					list = append(list, b.List[:i]...)
					list = append(list, newCells[:]...)
					list = append(list, b.List[i+1:]...)
					b.List = list
				}
			}
		}

		stop := &ast.ExprStmt{
		X: &ast.CallExpr{
			Fun:  &ast.SelectorExpr{ast.NewIdent(propnet),
					ast.NewIdent("Stop")},
			},
		}

		b.List = append(b.List, stop)
		fixed = true

	}, func(n interface{}) {
		b, ok := n.(*ast.Stmt)
		if !ok {
			return
		}

		*b = transformPropagator(*b)
	})

	return
}
